import torch as t
import torch.nn as nn
import time
import os


def weights_init(m):
    classname = m.__class__.__name__
    if classname.find('Conv') != -1:
        nn.init.normal_(m.weight.data, 0.0, 0.02)
    elif classname.find('BatchNorm') != -1:
        nn.init.normal_(m.weight.data, 1.0, 0.02)
        nn.init.constant_(m.bias.data, 0)

class BasicModule(t.nn.Module):
    """
    封装了nn.Module,主要是提供了save和load两个方法
    """
    def __init__(self):
        super(BasicModule,self).__init__()
        self.model_name = 'BasicModule'# 默认名字
        self.test = False

    def load(self, path, mached = True):
        """
        可加载指定路径的模型
        """
        model = t.load(path)
        self.load_state_dict(model, mached)
    
    #加载到cpu
    def loadCpu(self, path, mached = True):
        # map_location=torch.device('cpu')
        model = t.load(path, map_location=t.device('cpu'))
        self.load_state_dict(model, mached)
        
    #加载任意设备下的参数
    def loadAny(self, path, mached = True):
        print('loadany weight -> ', path)
        if next(self.parameters()).device.type == 'cpu':
            self.loadCpu(path, False)
        else:    
            self.load(path, False)

    def save(self, step):
        """
        保存模型，默认使用“模型名字+时间”作为文件名
        """
        prefix = 'weights/' + self.model_name + '_'
        if os.path.exists('weights/') == False:
            os.mkdir('weights/')
        name = '{}{}.pth'.format(prefix, step)
        t.save(self.state_dict(), name)
        return name
    
    #拷贝参数
    def copyWeights(self, model):
        self.load_state_dict(model.state_dict())

    def get_optimizer(self, lr, weight_decay):
        return t.optim.Adam(self.parameters(), lr=lr, weight_decay=weight_decay)
    
    #模型可视化
    def visualize(self, img):
        pass